import cv2
import time
from model import predict_image
from robotic_arm_controller import RoboticArmController
from camera_utils import initialize_camera, capture_frame, display_frame

if __name__ == "__main__":
    # 初始化机械臂控制器
    robotic_arm = RoboticArmController(port='COM4')  # 根据实际端口修改

    # 初始化摄像头
    cap = initialize_camera(1)  # 1表示外部摄像头

    # 状态变量
    last_label = None  # 上一次检测到的包装类型
    current_label = None  # 当前检测到的包装类型
    operation_delay = 8.0  # 操作延迟(秒)
    last_operation_time = time.time() - operation_delay  # 上次操作时间

    # 稳定检测计数器
    stable_count = 0
    stable_threshold = 8  # 连续8次相同结果才视为稳定

    # FPS计算
    fps = 0
    frame_count = 0
    start_time = time.time()

    # 操作状态
    operation_status = "待机"

    try:
        while True:
            # 捕获帧
            frame, pil_image = capture_frame(cap)
            if frame is None:
                break

            # 进行预测
            current_label = predict_image(pil_image)

            # 稳定检测机制
            if current_label == last_label:
                stable_count += 1
            else:
                stable_count = 0
                last_label = current_label

            # 计算时间差
            current_time = time.time()
            time_since_last = current_time - last_operation_time

            # 检测到稳定的包装类型且满足操作间隔
            if (stable_count >= stable_threshold and
                    time_since_last >= operation_delay and
                    not robotic_arm.operation_in_progress):

                print(f"检测到稳定的包装类型: {current_label}")
                operation_status = f"执行 {current_label} 操作"

                # 执行机械臂操作
                if robotic_arm.operate_for_package(current_label):
                    last_operation_time = time.time()
                    stable_count = 0  # 重置计数器
                    operation_status = "操作完成"
                else:
                    operation_status = "操作失败"

            # 计算FPS
            frame_count += 1
            elapsed_time = time.time() - start_time
            if elapsed_time >= 1.0:  # 每秒更新一次FPS
                fps = frame_count / elapsed_time
                start_time = time.time()
                frame_count = 0

            # 计算冷却时间
            cooldown = max(0, operation_delay - time_since_last)

            # 显示画面和信息
            display_frame(frame, current_label, cooldown, fps, operation_status)

            # 按'q'退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    finally:
        # 释放资源
        cap.release()
        cv2.destroyAllWindows()
        robotic_arm.close()
        print("程序已安全退出")